import numpy as np
import torch
import torch.optim as optim
import torch.nn.functional as F
from prepare_data import *
import argparse
import torch
import torch.nn as nn
from PIL import Image
from scipy.stats import pearsonr
from model import AesModel
from torch.utils.data import DataLoader
from torch.utils.data import Subset
from multiarmutils import *
from multiarmmodel import MultiArmedBanditGFS
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # # 加载小批量数据
    # train_set, test_set = load_aesthetic_data(opt.photo_size)
    # train_set = train_set[:100, :]  # 使用 100 条数据进行训练
    # test_set = test_set[:100, :]  # 使用 100 条数据进行测试

    train_setnew, test_setnew = get_RL_data(opt.photo_size)

    # 设置小型数据集的大小
    subset_size = 15  # 总样本数（批次数量 = subset_size / batch_size）

    # 构建子集
    train_subset = Subset(train_setnew, range(subset_size))
    test_subset = Subset(test_setnew, range(subset_size))

    train_data = DataLoader(dataset=train_subset, batch_size=5, collate_fn=collate_RL, pin_memory=True,
                            num_workers=0, drop_last=True)
    test_data = DataLoader(dataset=test_subset, batch_size=5, collate_fn=collate_RL, pin_memory=True,
                           num_workers=0, drop_last=True)

    # num_arms = train_set.shape[1]  # 每个美学维度作为一个“臂”
    exploration_ratio = 0.1  # 探索概率
    penalty_factor = 0.01
    # iterations = 10# 缩短迭代次数

    aes_model = AesModel(embedding_dim=opt.embedding_dim, opt=opt).to(device)
    aes_model.load_state_dict(torch.load('./aes_model_weight.pth'))


    initial_values = calculate_initial_values(train_data,aes_model)
    arms = [
        "幽默", "快乐", "好奇心", "享受和乐趣", "痛苦", "悲伤和厌倦", "审美愉悦", "有趣-无趣",
        "负面情绪", "不愉快-令人愉快", "无快乐-高度快乐", "惊喜", "放松-紧张", "观点", "视角",
        "视野", "理解", "共情", "共鸣", "同理心", "着迷", "感动", "感应", "积极的态度", "写作",
        "讲述更好的故事", "自发绘画", "线条", "颜色", "纹理", "大小", "插图", "文本", "图案", "色调",
        "边界", "形状", "亮度", "饱和度", "外观", "构图", "笔触", "形式", "内容", "留白", "注重细节的",
        "叠加", "线条绘制", "页面设计", "真实的", "自然的", "想象品质", "基于想象力的", "联想的观念",
        "颜色联想", "文化习得", "文化融合", "社会文化", "主题", "排版", "情节", "图像和语言的关联",
        "颜色关联", "整体结构", "简单-复杂", "清晰-不确定", "有序-无序", "不平衡-平衡", "类比",
        "艺术创作质量", "独特性", "艺术性", "可识别性", "意义", "有趣性", "复杂性", "充分性"
    ]

    # 初始化多臂赌博机
    mab = MultiArmedBanditGFS(
        # num_arms=num_arms,
        exploration_ratio=exploration_ratio,
        reward_function=calculate_total_reward,
        initial_values=initial_values,
        penalty_factor=penalty_factor,
        aes_model=aes_model,
        arms=arms
    )

    # 运行多臂赌博机
    final_selected_dims= mab.run(iterations=5, test_loader=test_data)
    # selected_dimensions = dynamic_annotation_selection(avg_dim_mse, threshold=0.02)
    print(f"选中的维度: {final_selected_dims}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--batchSize', type=int, default=10, help='input batch size')  # 原本是10000 现改5000
    parser.add_argument('--batch_size', type=int, default=20, help='input batch size')  # 原本是10000 现改5000
    parser.add_argument('--hidden_size', type=int, default=200, help='hidden state size')  # 原本是200 现改500
    parser.add_argument('--hidden_dim', type=int, default=100, help='hidden state size')  # 原本是200 现改500
    parser.add_argument('--epoch', type=int, default=1, help='number of epochs to train for')  # 这里为了方便调试代码，我改成1了
    parser.add_argument('--lr', type=float, default=0.1, help='learning rate')
    # parser.add_argument('--l2', type=float, default=0.0001, help='l2 penalty')
    parser.add_argument('--num_layers', type=int, default=6, help='layers nums')
    parser.add_argument('--num_heads', type=int, default=6, help='attention heads nums')
    parser.add_argument('--mlp_ratio', type=float, default=1, help='the ratio of hidden layers in the middle')
    parser.add_argument('--Kernel_size1', type=int, default=2, help='the first layer convolution kernel size')
    parser.add_argument('--Kernel_size2', type=int, default=2, help='the second layer convolution kernel size')
    parser.add_argument('--Stride1', type=int, default=2, help='the second layer convolution stride size')
    parser.add_argument('--Stride2', type=int, default=2, help='the second layer convolution stride size')
    parser.add_argument('--num_classes', type=int, default=77, help='the number of categories')
    parser.add_argument('--num_classes_last', type=int, default=7, help='the number of categories to the last label')
    parser.add_argument('--photo_size', type=int, default=128, help='the number of categories to the last label')
    parser.add_argument('--Linear_nums', type=int, default=3, help='the number of categories to the last label')
    parser.add_argument('--pb_path', type=str, default='./data', help='the number of categories to the last label')
    parser.add_argument('--state_nums', type=int, default=77, help='the number of categories to the last label')
    parser.add_argument('--gamma', type=float, default=0.99, help='the number of categories to the last label')
    parser.add_argument('--epsilon', type=float, default=0.8, help='the number of categories to the last label')
    parser.add_argument('--target_update_nums', type=int, default=5, help='the number of categories to the last label')
    parser.add_argument('--least_score', type=float, default=0.8, help='the number of categories to the last label')
    parser.add_argument('--ReplayBuffer_capacity', type=int, default=100,
                        help='the number of categories to the last label')
    parser.add_argument('--min_size', type=int, default=50, help='the number of categories to the last label')
    parser.add_argument('--path_len', type=int, default=5, help='the number of categories to the last label')
    parser.add_argument('--D', type=float, default=1.7, help='the number of categories to the last label')
    parser.add_argument('--a', type=int, default=5, help='the number of categories to the last label')
    parser.add_argument('--model_name', type=str, default=str('Aes'), help='the number of categories to the last label')
    parser.add_argument('--mu', type=float, default=1, help='the number of categories to the last label')
    parser.add_argument('--embedding_dim', type=int, default=192, help='the number of categories to the last label')
    parser.add_argument('--msepara', type=int, default=100, help='the number of categories to the last label')
    parser.add_argument('--expand_name', type=str, default=str('embedding_dim'),
                        help='the number of categories to the last label')
    parser.add_argument('--is_adjust_parameter', type=str, default=str('true'),
                        help='the number of categories to the last label')
    parser.add_argument("--mse_threshold", type=float, default=0.01, help="MSE阈值，用于调整top_k")
    opt = parser.parse_args()
    main()